import gradio as gr
import os
import json
import time
import shutil
import uuid
from typing import Optional, Tuple

from Core.comfy_client import ComfyUIClient
from Core.baidu_translate import baidu_translate


class FillLoRAInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        self.workflow_path = os.path.join(
            os.path.dirname(__file__), "..", "Comfyui", "work_flow", "Flux_fill_lora.json"
        )

        # 目录：unet、clip、loras、input、temp
        from config.config import get_unet_dir_path, get_clip_dir_path, get_lora_dir_path, get_input_path
        self.unet_dir = get_unet_dir_path()
        self.clip_dir = get_clip_dir_path()
        self.lora_dir = get_lora_dir_path()
        self.input_dir = get_input_path()
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)
        os.makedirs(self.input_dir, exist_ok=True)

        # 模型列表
        self.unet_files = [f for f in os.listdir(self.unet_dir)] if os.path.exists(self.unet_dir) else []
        self.unet_files = [f for f in self.unet_files if f.lower().endswith((".safetensors", ".ckpt"))]
        self.clip_files = [f for f in os.listdir(self.clip_dir)] if os.path.exists(self.clip_dir) else []
        self.clip_files = [f for f in self.clip_files if f.lower().endswith((".safetensors", ".ckpt"))]
        self.lora_choices = self._collect_lora_choices()

        # 默认值
        self.default_unet = "flux1-dev-fp8.safetensors"
        if self.unet_files and self.default_unet not in self.unet_files:
            self.default_unet = self.unet_files[0]
        self.default_clip2 = "t5xxl_fp8_e4m3fn.safetensors"
        if self.clip_files and self.default_clip2 not in self.clip_files:
            self.default_clip2 = self.clip_files[0]
        self.default_lora = "test\\ETShoeF.safetensors"
        if self.lora_choices and self.default_lora not in self.lora_choices:
            self.default_lora = self.lora_choices[0]

    def _collect_lora_choices(self):
        choices = []
        try:
            if os.path.exists(self.lora_dir):
                for root, _, files in os.walk(self.lora_dir):
                    for f in files:
                        if f.lower().endswith((".safetensors", ".ckpt")):
                            rel = os.path.relpath(os.path.join(root, f), self.lora_dir)
                            rel = rel.replace("/", "\\")
                            choices.append(rel)
        except Exception as e:
            print(f"收集LoRA失败: {e}")
        choices.sort()
        return choices

    def load_workflow(self):
        try:
            with open(self.workflow_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"加载工作流失败: {e}")
            return None

    def _copy_image_to_input(self, src_path: str, prefix: str) -> Tuple[Optional[str], Optional[str]]:
        try:
            if not src_path or not os.path.exists(src_path):
                return None, "错误：未提供有效的图像路径"
            ext = os.path.splitext(src_path)[1].lower()
            if ext not in [".png", ".jpg", ".jpeg", ".bmp", ".webp", ".tiff"]:
                return None, "错误：仅支持PNG/JPG/JPEG/BMP/WEBP/TIFF"
            filename = f"{prefix}_{uuid.uuid4().hex}{ext}"
            dst = os.path.join(self.input_dir, filename)
            shutil.copy2(src_path, dst)
            return filename, None
        except Exception as e:
            return None, f"复制图片失败：{e}"

    def _save_mask_to_temp(self, mask_data) -> Optional[str]:
        """参考 Tab_Redux_ReduxMask 的简化保存：将蒙版保存为PNG到temp并返回文件路径。"""
        try:
            import numpy as np
            from PIL import Image
            import cv2
            mask_path = os.path.join(self.temp_dir, f"mask_{uuid.uuid4().hex}.png")
            if mask_data is None:
                return None
            if isinstance(mask_data, np.ndarray):
                Image.fromarray(mask_data).save(mask_path)
                return mask_path
            if isinstance(mask_data, dict) and "layers" in mask_data and len(mask_data["layers"]) > 0:
                cv2.imwrite(mask_path, mask_data["layers"][0])
                return mask_path
            return None
        except Exception as e:
            print(f"保存蒙版失败: {e}")
            return None

    def _save_masked_alpha_png(self, base_image_path, mask_data) -> Tuple[Optional[str], Optional[str]]:
        """将蒙版合成到原图的 Alpha 通道并保存到 input。
        规则：ImageEditor 黑色笔触区域视为需要重绘区域；为满足 InpaintModelConditioning，
        将重绘区域的 alpha 置为 0（透明），非重绘区域 alpha 置为 255（不透明）。"""
        try:
            if not base_image_path or not os.path.exists(base_image_path):
                return None, "错误：未提供有效的基础图像路径"
            from PIL import Image
            import numpy as np

            base = Image.open(base_image_path).convert("RGBA")

            # 提取编辑器蒙版：优先使用 layers[0] 的 alpha；其次使用 numpy 数组推断
            mask_alpha = None
            try:
                if isinstance(mask_data, dict):
                    # 有些情况下编辑器会提供完整 RGBA 层
                    if "layers" in mask_data and len(mask_data["layers"]) > 0:
                        m = mask_data["layers"][0]
                        if isinstance(m, np.ndarray):
                            if m.ndim == 3:
                                # 使用图层的 alpha，alpha>0 表示绘制区域
                                if m.shape[2] >= 4:
                                    mask_alpha = m[:, :, 3]
                                else:
                                    # 无 alpha 时，取最大通道作为近似蒙版
                                    mask_alpha = m.max(axis=2)
                            elif m.ndim == 2:
                                mask_alpha = m
                    # 某些返回结构可能包含直接的 mask 数组
                    elif "mask" in mask_data and isinstance(mask_data["mask"], np.ndarray):
                        m = mask_data["mask"]
                        if m.ndim == 3:
                            mask_alpha = m[:, :, 3] if m.shape[2] >= 4 else m.max(axis=2)
                        elif m.ndim == 2:
                            mask_alpha = m
                elif isinstance(mask_data, np.ndarray):
                    m = mask_data
                    if m.ndim == 3:
                        mask_alpha = m[:, :, 3] if m.shape[2] >= 4 else m.max(axis=2)
                    elif m.ndim == 2:
                        mask_alpha = m
            except Exception:
                mask_alpha = None

            # 构建最终 alpha：将“绘制区域”视为需要重绘 -> alpha=0；其余 alpha=255
            if mask_alpha is None:
                final_alpha = Image.new("L", base.size, 255)
            else:
                # 归一化为 0/255 的 uint8
                mask_alpha = mask_alpha.astype(np.float32)
                mmin, mmax = mask_alpha.min(), mask_alpha.max()
                if mmax > mmin:
                    mask_alpha = (mask_alpha - mmin) / (mmax - mmin)
                mask_alpha = (mask_alpha * 255.0).clip(0, 255).astype(np.uint8)

                # 绘制区域判断：alpha>0 为绘制痕迹，将其转为 0（透明）
                drawn = (mask_alpha > 0).astype(np.uint8) * 255
                inv_alpha = 255 - drawn  # 绘制区域 -> 0；非绘制 -> 255

                final_alpha = Image.fromarray(inv_alpha)
                if final_alpha.size != base.size:
                    final_alpha = final_alpha.resize(base.size)

                # 如果用户未绘制（整张 inv_alpha=255），保持全不透明，避免空蒙版导致工作流异常
                if np.sum(np.array(final_alpha)) == 255 * final_alpha.size[0] * final_alpha.size[1]:
                    final_alpha = Image.new("L", base.size, 255)

            r, g, b, _ = base.split()
            out = Image.merge("RGBA", (r, g, b, final_alpha))

            filename = f"inpaint_masked_{uuid.uuid4().hex}.png"
            dst = os.path.join(self.input_dir, filename)
            out.save(dst)
            return filename, None
        except Exception as e:
            return None, f"合成蒙版失败：{e}"

    def generate_image(
        self,
        lora_ref_image_path,
        inpaint_image_path,
        inpaint_mask_data,
        unet_name,
        clip_name2,
        lora_name,
        strength_model,
        steps,
        denoise,
        text_input,
        max_retries: int = 40,
    ):
        try:
            workflow = self.load_workflow()
            if not workflow:
                return None, "无法加载工作流", None

            # 文本翻译（节点39 TextInput_）
            try:
                english = baidu_translate(text_input or "")
                translate_info = "翻译成功"
            except Exception as e:
                english = text_input or ""
                translate_info = f"翻译失败，改用原文本: {str(e)}"
            if "39" in workflow and "inputs" in workflow["39"]:
                workflow["39"]["inputs"]["text"] = english

            # 节点1：UNET模型
            if "1" in workflow and "inputs" in workflow["1"]:
                workflow["1"]["inputs"]["unet_name"] = unet_name

            # 节点8：DualCLIP clip_name2
            if "8" in workflow and "inputs" in workflow["8"]:
                workflow["8"]["inputs"]["clip_name2"] = clip_name2

            # 节点41：LoRA加载器
            if "41" in workflow and "inputs" in workflow["41"]:
                workflow["41"]["inputs"]["lora_name"] = lora_name
                workflow["41"]["inputs"]["strength_model"] = float(strength_model)

            # 节点11：KSampler步数与降噪
            if "11" in workflow and "inputs" in workflow["11"]:
                workflow["11"]["inputs"]["steps"] = int(steps)
                workflow["11"]["inputs"]["denoise"] = float(denoise)

            # 节点56：LoRA参考图（复制到input目录）
            ref_filename, err_ref = self._copy_image_to_input(lora_ref_image_path, "lora_ref")
            if err_ref:
                return None, f"生成失败：{err_ref}", english
            if "56" in workflow and "inputs" in workflow["56"]:
                workflow["56"]["inputs"]["image"] = ref_filename

            # 节点16：局部重绘图片 + 蒙版处理
            # 1) 保存一份蒙版到temp（参考ReduxMask简化逻辑）
            _ = self._save_mask_to_temp(inpaint_mask_data)
            # 2) 将蒙版融合为Alpha并保存到input，供LoadImage输出mask到节点12使用
            inpaint_masked_filename, err_inp = self._save_masked_alpha_png(inpaint_image_path, inpaint_mask_data)
            if err_inp:
                return None, f"生成失败：{err_inp}", english
            if "16" in workflow and "inputs" in workflow["16"]:
                workflow["16"]["inputs"]["image"] = inpaint_masked_filename

            # 提交生成请求
            response = self.client.post_prompt(workflow)
            prompt_id = response.get("prompt_id")
            if not prompt_id:
                return None, "生成失败：未返回prompt_id", english

            # 轮询获取结果图
            retry_interval = 2
            for _ in range(max_retries):
                history = self.client.get_history(prompt_id)
                if prompt_id in history and "outputs" in history[prompt_id]:
                    outputs = history[prompt_id]["outputs"]
                    for node_id, data in outputs.items():
                        if "images" in data:
                            for image in data["images"]:
                                img_bytes = self.client.get_image(image["filename"], image.get("subfolder"), image.get("type"))
                                result_file = os.path.join(self.temp_dir, f"final_{prompt_id}.png")
                                with open(result_file, "wb") as f:
                                    f.write(img_bytes)
                                # 同步保存到项目 output_images
                                try:
                                    from config.config import get_local_output_images_path
                                    out_dir = get_local_output_images_path()
                                    os.makedirs(out_dir, exist_ok=True)
                                    shutil.copy2(result_file, os.path.join(out_dir, os.path.basename(result_file)))
                                except Exception as e:
                                    print(f"复制到output_images失败: {e}")
                                status = "生成成功\n" + translate_info
                                return os.path.abspath(result_file), status, english
                time.sleep(retry_interval)

            return None, "生成失败：超时或无输出", english
        except Exception as e:
            print(f"生成异常: {e}")
            return None, f"生成失败：{e}", None


def Tab_fill_LoRA():
    interface = FillLoRAInterface()

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 加载LoRA局部重新设计")

            # 节点56：LoRA参考图
            lora_ref_image = gr.Image(label="上传LoRA参考图", type="filepath", interactive=True)

            # 节点16：局部重绘图片 + 蒙版绘制（黑色画笔）
            inpaint_image = gr.Image(label="上传需要局部重绘的图片", type="filepath", interactive=True)
            inpaint_editor = gr.ImageEditor(
                label="绘制重绘区域（黑色）",
                type="numpy",
                brush=gr.Brush(colors=["#000000"], color_mode="fixed"),
                interactive=True,
            )

            def update_inpaint_editor(img_path):
                return img_path

            inpaint_image.change(fn=update_inpaint_editor, inputs=inpaint_image, outputs=inpaint_editor)

            # 节点1：UNET选择
            unet_dropdown = gr.Dropdown(
                label="UNET模型(unet_name)",
                choices=interface.unet_files,
                value=interface.default_unet,
                interactive=bool(interface.unet_files),
            )

            # 节点8：clip_name2 模型
            clip2_dropdown = gr.Dropdown(
                label="CLIP文本编码器(clip_name2)",
                choices=interface.clip_files,
                value=interface.default_clip2,
                interactive=bool(interface.clip_files),
            )

            # 节点41：LoRA加载器
            lora_dropdown = gr.Dropdown(
                label="LoRA名称(lora_name)",
                choices=interface.lora_choices,
                value=interface.default_lora,
                interactive=bool(interface.lora_choices),
            )
            strength_model_slider = gr.Slider(label="LoRA模型强度(strength_model)", minimum=0, maximum=1, value=0.8, step=0.01)

            # 节点11：KSampler
            steps_input = gr.Number(label="步数(steps)", value=20, minimum=0, maximum=60)
            denoise_slider = gr.Slider(label="降噪(denoise)", minimum=0, maximum=1, value=1, step=0.01)

            # 节点39：文本 + 翻译展示
            text_input = gr.Textbox(label="文本(中文自动翻译)", lines=3)
            translated_text = gr.Textbox(label="英文翻译(用于工作流)", lines=3, interactive=False)

        with gr.Column(scale=2):
            generate_btn = gr.Button("生成", variant="primary")
            result_image = gr.Image(label="生成结果")
            status_text = gr.Textbox(label="状态", interactive=False)

            def get_output_images():
                from config.config import get_local_output_images_path
                out_dir = get_local_output_images_path()
                if not os.path.exists(out_dir):
                    return []
                images = []
                for file in os.listdir(out_dir):
                    if file.lower().endswith((".png", ".jpg", ".jpeg")):
                        images.append(os.path.join(out_dir, file))
                images.sort(key=lambda x: os.path.getmtime(x), reverse=True)
                return images

            with gr.Row():
                refresh_btn = gr.Button("刷新输出预览")
            with gr.Row():
                output_gallery = gr.Gallery(label="输出文件夹图片", columns=4, show_label=True, elem_id="output_gallery_fill_lora")

            refresh_btn.click(fn=get_output_images, outputs=[output_gallery])

    def translate_and_generate(lora_ref_path, inpaint_img_path, inpaint_mask_data, unet_name, clip2_name, lora_name, strength_model, steps, denoise, text):
        img_path, status, english = interface.generate_image(
            lora_ref_path,
            inpaint_img_path,
            inpaint_mask_data,
            unet_name,
            clip2_name,
            lora_name,
            strength_model,
            steps,
            denoise,
            text,
        )
        return img_path, status, english

    generate_btn.click(
        fn=translate_and_generate,
        inputs=[
            lora_ref_image,
            inpaint_image,
            inpaint_editor,
            unet_dropdown,
            clip2_dropdown,
            lora_dropdown,
            strength_model_slider,
            steps_input,
            denoise_slider,
            text_input,
        ],
        outputs=[result_image, status_text, translated_text],
    )